{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "### Seemore: Vision Language Model from scratch in Pytorch\n",
        "\n",
        "TL;DR: In  this blog I implement a vision language model consisting of an image encoder, a multimodal projection module and a decoder language model in pure pytorch. The name ‘seemore’ is my way of paying homage to Andrej Karpathy’s project ‘makemore’ because here I use a character level autoregressive language model much like in his nanoGPT/ makemore implementation.  My goal is for you to have an intuitive understanding of how it all works once you read this blog and step through the code in the repo.\n",
        "\n",
        "\n",
        "  <img src=\"https://github.com/AviSoori1x/seemore/blob/main/images/seemorelogo.png?raw=true\" width=\"300\" height=\"300\" alt=\"seemore\">\n",
        "\n"
      ],
      "metadata": {
        "id": "12xieaXiYZBn"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In ‘seemore’, my simple implementation of a vision language model (VLM), there are 3 main components.\n",
        "\n",
        "  <img src=\"https://github.com/AviSoori1x/seemore/blob/main/images/vlm.png?raw=true\" width=\"500\" height=\"500\" alt=\"seemore\">\n",
        "\n",
        "* Image Encoder to extract visual features from images. In this case I use a from scratch implementation of the original vision transformer used in CLIP. This is actually a popular choice in many modern VLMs. The one notable exception is Fuyu series of models from Adept, that passes the patchified images directly to the projection layer.\n",
        "\n",
        "* Vision-Language Projector - Image embeddings are not of the same shape as text embeddings used by the decoder. So we need to ‘project’ i.e. change dimensionality of image features extracted by the image encoder to match what’s observed in the text embedding space. So image features become ‘visual tokens’ for the decoder. This could be a single layer or an MLP. I’ve used an MLP because it’s worth showing.\n",
        "\n",
        "* A decoder only language model. This is the component that ultimately generates text. In my implementation I’ve deviated from what you see in LLaVA a bit by incorporating the projection module to my decoder. Typically this is not observed, and you leave the architecture of the decoder (which is usually an already pretrained model) untouched.\n",
        "\n"
      ],
      "metadata": {
        "id": "mhbgAfWAZbqj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Dop the necessary imports. Note that this is just standard pytorch"
      ],
      "metadata": {
        "id": "btPeznnpZQIl"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "7423d594-7c2c-4a5b-8089-3faf9e30b116",
          "showTitle": false,
          "title": ""
        },
        "id": "Ozzn44gSRw5F"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "import io\n",
        "import pandas as pd\n",
        "from PIL import Image\n",
        "import torchvision.transforms as transforms\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import functional as F\n",
        "from torch.nn import init"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "71cb328e-20a2-42e7-95d2-ccdf9bcc238a",
          "showTitle": false,
          "title": ""
        },
        "id": "hcu64FfMRw5H"
      },
      "outputs": [],
      "source": [
        "# Ensure every computation happens on the GPU when available\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "3e279fd6-d26a-481e-b428-8629635da103",
          "showTitle": false,
          "title": ""
        },
        "id": "d2LDm8CmRw5I"
      },
      "outputs": [],
      "source": [
        "#To build the encoding and decoding functions we use the tinyshakespear dataset. However for the sake of brevity we do not pretrain the decoder model on it\n",
        "#the training function should be able to do it without an issue as well as it could take both images and tex\n",
        "text_path = \"./input.txt\"\n",
        "with open(text_path, 'r', encoding='utf-8') as f:\n",
        "    text = f.read()\n",
        "\n",
        "# here are all the unique characters that occur in this text\n",
        "chars = sorted(list(set(text)))\n",
        "# create a mapping from characters to integers\n",
        "stoi = { ch:i for i,ch in enumerate(chars) }\n",
        "stoi['<pad>']= 65\n",
        "itos = { i:ch for i,ch in enumerate(chars) }\n",
        "itos[65] = '<pad>'\n",
        "encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list of integers\n",
        "decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string\n",
        "vocab_size = len(stoi.keys())"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "<img src=\"https://github.com/AviSoori1x/seemore/blob/main/images/clip.png?raw=true\" width=\"600\" height=\"400\" alt=\"seemore\">\n",
        "\n",
        "Usually a pretrained vision transformer from CLIP or SigLIP (an improved version of CLIP) is used.\n",
        "\n"
      ],
      "metadata": {
        "id": "OwKjyYG_k9n8"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "<img src=\"https://github.com/AviSoori1x/seemore/blob/main/images/vit.png?raw=true\" width=\"600\" height=\"300\" alt=\"seemore\">\n",
        "\n",
        "To implement this vision transformer from scratch we have to create a PatchEmbeddings class that can take an image and create a sequence of patches. This process is crucial for enabling the transformer architecture to process visual data effectively, specifically using the attention blocks in the subsequent steps of the architecture.\n"
      ],
      "metadata": {
        "id": "rQx-syMneYi8"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "08daa24b-a3e9-46f9-bc57-51d8d1673397",
          "showTitle": false,
          "title": ""
        },
        "id": "TH449G_WRw5I"
      },
      "outputs": [],
      "source": [
        "class PatchEmbeddings(nn.Module):\n",
        "    def __init__(self, img_size=96, patch_size=16, hidden_dim=512):\n",
        "        super().__init__()\n",
        "\n",
        "        # Store the input image size\n",
        "        self.img_size = img_size\n",
        "\n",
        "        # Store the size of each patch\n",
        "        self.patch_size = patch_size\n",
        "\n",
        "        # Calculate the total number of patches\n",
        "        self.num_patches = (img_size // patch_size) ** 2\n",
        "\n",
        "        # Create a convolutional layer to extract patch embeddings\n",
        "        # in_channels=3 assumes the input image has 3 color channels (RGB)\n",
        "        # out_channels=hidden_dim sets the number of output channels to match the hidden dimension\n",
        "        # kernel_size=patch_size and stride=patch_size ensure each patch is separately embedded\n",
        "        self.conv = nn.Conv2d(in_channels=3, out_channels=hidden_dim,\n",
        "                              kernel_size=patch_size, stride=patch_size)\n",
        "\n",
        "    def forward(self, X):\n",
        "        # Extract patch embeddings from the input image\n",
        "        X = self.conv(X)\n",
        "\n",
        "        # Flatten the spatial dimensions (height and width) of the patch embeddings\n",
        "        # This step flattens the patch dimensions into a single dimension\n",
        "        X = X.flatten(2)\n",
        "\n",
        "        # Transpose the dimensions to obtain the shape [batch_size, num_patches, hidden_dim]\n",
        "        # This step brings the num_patches dimension to the second position\n",
        "        X = X.transpose(1, 2)\n",
        "\n",
        "        return X\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "5b3934cf-5735-470c-aece-7babbd097a38",
          "showTitle": false,
          "title": ""
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mhTi7MpFRw5I",
        "outputId": "f9c27e9f-a38b-49ce-a5a7-b426fa999bab"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([4, 36, 512])"
            ]
          },
          "metadata": {},
          "execution_count": 18
        }
      ],
      "source": [
        "#testing\n",
        "img_size, patch_size,  num_hiddens, batch_size = 96, 16, 512, 4\n",
        "patch_embeddings = PatchEmbeddings(img_size, patch_size, num_hiddens )\n",
        "X = torch.zeros(batch_size, 3, img_size, img_size)\n",
        "patch_embeddings(X).shape"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Things get interesting when building the components seen in the transformer blocks. I.e. The attention head implementation, multi head attention, the multilayer perceptron seen in each attention head and the transformer block itself. These components are mostly identical across the vision transformer we are implementing for the ‘visual token’ generation and the decoder language model for the actual text output generation.\n",
        "\n",
        "The only key difference is the masking applied in each attention head in the decoder language model. This is done to ensure the integrity of the autoregressive language generation process, particularly in a decoder-only model, the code implements masking. This masking technique is crucial as it obscures any information following the current token's position, thereby directing the model's attention to only the preceding parts of the sequence. Such an attention mechanism is known as causal self-attention.\n"
      ],
      "metadata": {
        "id": "EtQiv8I7fQMW"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The multilayer perceptron that follows each multihead attention module is quite straightforward. Please note that I’ve noticed GELU used quite often in Vision Transformers and ReLU used in text transformers, so I have this conditional logic to switch between the two based on where this MLP will be inserted. However, it seems that GELU is being used for both due to its resultant model performance, regardless of the fact that it’s more computationally expensive that RELU.\n"
      ],
      "metadata": {
        "id": "mEN3qhAZfd6S"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "acec0b6d-4b53-497a-b0a4-8e8aee8270ee",
          "showTitle": false,
          "title": ""
        },
        "id": "_xJdDCXgRw5I"
      },
      "outputs": [],
      "source": [
        "#swapping linear for lazy linear for simplicity. Lazylinear can accept any arbitrary input dimension without having it specified\n",
        "\n",
        "class MLP(nn.Module):\n",
        "    def __init__(self, n_embd, dropout=0.1, is_decoder=True):\n",
        "        super().__init__()\n",
        "\n",
        "        # Define the layers of the MLP\n",
        "        layers = [\n",
        "            # First linear layer that expands the input dimension from n_embd to 4 * n_embd\n",
        "            nn.Linear(n_embd, 4 * n_embd),\n",
        "\n",
        "            # Activation function: ReLU if is_decoder is True, else GELU\n",
        "            nn.ReLU() if is_decoder else nn.GELU(),\n",
        "\n",
        "            # Second linear layer that projects the intermediate dimension back to n_embd\n",
        "            nn.Linear(4 * n_embd, n_embd),\n",
        "\n",
        "            # Dropout layer for regularization\n",
        "            nn.Dropout(dropout)\n",
        "        ]\n",
        "\n",
        "        # Create a sequential container to hold the layers\n",
        "        self.net = nn.Sequential(*layers)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Pass the input through the MLP layers\n",
        "        return self.net(x)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "e56112ea-2647-45c2-b9ca-12f78a6827c7",
          "showTitle": false,
          "title": ""
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ct1pFoFpRw5J",
        "outputId": "13c3618b-e471-4822-8350-9c767c7bf31a"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([4, 3, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ],
      "source": [
        "#For the sake of this example consider embedding size to be 128\n",
        "n_embd = 128\n",
        "testmlp = MLP(n_embd)\n",
        "mlp_input = torch.zeros(batch_size, 3, n_embd)\n",
        "testmlp_out = testmlp(mlp_input)\n",
        "testmlp_out.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "When it comes to the self attention mechanism employed across both the vision encoder and the language decoder, the only key difference is the masking applied in each attention head in the decoder language model. This is done to ensure the integrity of the autoregressive language generation process, particularly in a decoder-only model, the code implements masking. This masking technique is crucial as it obscures any information following the current token's position, thereby directing the model's attention to only the preceding parts of the sequence. Such an attention mechanism is known as causal self-attention.\n",
        "\n",
        "<img src=\"https://github.com/AviSoori1x/seemore/blob/main/images/mhsa.png?raw=true\" width=\"800\" height=\"400\" alt=\"seemore\">\n",
        "\n",
        "In the above image, the lower triangular mask is only applied in the case of a decoder model. Consider the bright blue triangle in matrix W absent in the case of visualizing the process in each attention head in the vision encoder.\n"
      ],
      "metadata": {
        "id": "AurTlZV7g3ov"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "fff7d282-c2ab-4b90-ab4a-f99fa9cd14c8",
          "showTitle": false,
          "title": ""
        },
        "id": "uqj-YL7pRw5J"
      },
      "outputs": [],
      "source": [
        "class Head(nn.Module):\n",
        "    def __init__(self, n_embd, head_size, dropout=0.1, is_decoder=False):\n",
        "        super().__init__()\n",
        "\n",
        "        # Linear layer for key projection\n",
        "        self.key = nn.Linear(n_embd, head_size, bias=False)\n",
        "\n",
        "        # Linear layer for query projection\n",
        "        self.query = nn.Linear(n_embd, head_size, bias=False)\n",
        "\n",
        "        # Linear layer for value projection\n",
        "        self.value = nn.Linear(n_embd, head_size, bias=False)\n",
        "\n",
        "        # Dropout layer for regularization\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "        # Flag indicating whether this head is used in the decoder\n",
        "        self.is_decoder = is_decoder\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Get the batch size (B), sequence length (T), and embedding dimension (C) from the input tensor\n",
        "        B, T, C = x.shape\n",
        "\n",
        "        # Compute key, query, and value projections\n",
        "        k = self.key(x)   # Shape: [B, T, head_size]\n",
        "        q = self.query(x) # Shape: [B, T, head_size]\n",
        "        v = self.value(x) # Shape: [B, T, head_size]\n",
        "\n",
        "        # Compute attention scores by taking the dot product of query and key\n",
        "        # and scaling by the square root of the embedding dimension\n",
        "        wei = q @ k.transpose(-2, -1) * (C ** -0.5) # Shape: [B, T, T]\n",
        "\n",
        "        if self.is_decoder:\n",
        "            # If this head is used in the decoder, apply a causal mask to the attention scores\n",
        "            # to prevent attending to future positions\n",
        "            tril = torch.tril(torch.ones(T, T, dtype=torch.bool, device=x.device))\n",
        "            wei = wei.masked_fill(tril == 0, float('-inf'))\n",
        "\n",
        "        # Apply softmax to the attention scores to obtain attention probabilities\n",
        "        wei = F.softmax(wei, dim=-1) # Shape: [B, T, T]\n",
        "\n",
        "        # Apply dropout to the attention probabilities for regularization\n",
        "        wei = self.dropout(wei)\n",
        "\n",
        "        # Perform weighted aggregation of values using the attention probabilities\n",
        "        out = wei @ v # Shape: [B, T, head_size]\n",
        "\n",
        "        return out\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "17f9b16f-0ebc-4b59-b285-78790613aa15",
          "showTitle": false,
          "title": ""
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FmD_ct_-Rw5J",
        "outputId": "9bf43428-b2c0-41b5-99c1-32955332f0ca"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([4, 3, 16])"
            ]
          },
          "metadata": {},
          "execution_count": 22
        }
      ],
      "source": [
        "#Example values for testing\n",
        "n_embd, head_size, batch_size = 128, 16, 4\n",
        "\n",
        "testhead = Head(n_embd, head_size)\n",
        "head_input = torch.zeros(batch_size, 3, n_embd)\n",
        "testhead_out = testhead(head_input)\n",
        "testhead_out.shape # (B, T,H_size)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "49de410a-f93b-4413-a06e-a7487379c5fc",
          "showTitle": false,
          "title": ""
        },
        "id": "-xHZhwnHRw5J"
      },
      "outputs": [],
      "source": [
        "class MultiHeadAttention(nn.Module):\n",
        "    def __init__(self, n_embd, num_heads, dropout=0.1, is_decoder=False):\n",
        "        super().__init__()\n",
        "\n",
        "        # Ensure that the embedding dimension is divisible by the number of heads\n",
        "        assert n_embd % num_heads == 0, \"n_embd must be divisible by num_heads\"\n",
        "\n",
        "        # Create a ModuleList of attention heads\n",
        "        self.heads = nn.ModuleList([\n",
        "            Head(n_embd, n_embd // num_heads, dropout, is_decoder)\n",
        "            for _ in range(num_heads)\n",
        "        ])\n",
        "\n",
        "        # Linear layer for projecting the concatenated head outputs\n",
        "        self.proj = nn.Linear(n_embd, n_embd)\n",
        "\n",
        "        # Dropout layer for regularization\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Apply each attention head to the input tensor\n",
        "        head_outputs = [h(x) for h in self.heads]\n",
        "\n",
        "        # Concatenate the outputs from all heads along the last dimension\n",
        "        out = torch.cat(head_outputs, dim=-1)\n",
        "\n",
        "        # Apply the projection layer to the concatenated outputs\n",
        "        out = self.proj(out)\n",
        "\n",
        "        # Apply dropout to the projected outputs for regularization\n",
        "        out = self.dropout(out)\n",
        "\n",
        "        return out\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "90c71fe2-3ef6-4eca-9805-130b12418646",
          "showTitle": false,
          "title": ""
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qudiJJl7Rw5J",
        "outputId": "0ace91c4-e682-4540-beec-ac18b742cf25"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([4, 3, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 24
        }
      ],
      "source": [
        "#Example values for testing\n",
        "n_embd, n_head = 128, 8\n",
        "testmha = MultiHeadAttention(n_embd, n_head)\n",
        "head_input = torch.zeros(batch_size, 3, n_embd)\n",
        "testmha_out = testmha(head_input)\n",
        "testmha_out.shape # (B, T,H_size*n_heads = n_embed)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Each encoder transformer block looks as follows"
      ],
      "metadata": {
        "id": "GCesbi9PiFe_"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "eb477c79-fd05-4e79-8149-27b62d2b1c7e",
          "showTitle": false,
          "title": ""
        },
        "id": "xmSrzvZYRw5K"
      },
      "outputs": [],
      "source": [
        "class Block(nn.Module):\n",
        "    def __init__(self, n_embd, num_heads, dropout=0.1, is_decoder=False):\n",
        "        super().__init__()\n",
        "\n",
        "        # Layer normalization for the input to the attention layer\n",
        "        self.ln1 = nn.LayerNorm(n_embd)\n",
        "\n",
        "        # Multi-head attention module\n",
        "        self.attn = MultiHeadAttention(n_embd, num_heads, dropout, is_decoder)\n",
        "\n",
        "        # Layer normalization for the input to the FFN\n",
        "        self.ln2 = nn.LayerNorm(n_embd)\n",
        "\n",
        "        # Feed-forward neural network (FFN)\n",
        "        self.ffn = nn.Sequential(\n",
        "            nn.Linear(n_embd, 4 * n_embd),  # Expand the dimension\n",
        "            nn.GELU(),  # Activation function\n",
        "            nn.Linear(4 * n_embd, n_embd),  # Project back to the original dimension\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        original_x = x  # Save the input for the residual connection\n",
        "\n",
        "        # Apply layer normalization to the input\n",
        "        x = self.ln1(x)\n",
        "\n",
        "        # Apply multi-head attention\n",
        "        attn_output = self.attn(x)\n",
        "\n",
        "        # Add the residual connection (original input) to the attention output\n",
        "        x = original_x + attn_output\n",
        "\n",
        "        # Apply layer normalization to the input to the FFN\n",
        "        x = self.ln2(x)\n",
        "\n",
        "        # Apply the FFN\n",
        "        ffn_output = self.ffn(x)\n",
        "\n",
        "        # Add the residual connection (input to FFN) to the FFN output\n",
        "        x = x + ffn_output\n",
        "\n",
        "        return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "ed52337a-c5ff-4164-bb5d-0e450c3a02fd",
          "showTitle": false,
          "title": ""
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NBxryt6vRw5K",
        "outputId": "f8954d5c-ce75-4e99-95c2-034ef3c0f1d3"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([4, 3, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 26
        }
      ],
      "source": [
        "#Example values for testing\n",
        "n_embd, head_size, batch_size = 128, 16, 4\n",
        "\n",
        "testblock = Block(n_embd, n_head)\n",
        "block_input = torch.zeros(batch_size, 3, n_embd)\n",
        "testblock_out = testblock(block_input)\n",
        "testblock_out.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now all this can be be put together to implement a Vision Transformer"
      ],
      "metadata": {
        "id": "fRs4ppwaia9z"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "fd6bdd3c-3c0c-413b-8800-a036a7848492",
          "showTitle": false,
          "title": ""
        },
        "id": "bSBuAN1GRw5K"
      },
      "outputs": [],
      "source": [
        "class ViT(nn.Module):\n",
        "    def __init__(self, img_size, patch_size, num_hiddens, num_heads, num_blks, emb_dropout, blk_dropout):\n",
        "        super().__init__()\n",
        "\n",
        "        # Patch embedding layer to convert the input image into patches\n",
        "        self.patch_embedding = PatchEmbeddings(img_size, patch_size, num_hiddens)\n",
        "\n",
        "        # Learnable classification token\n",
        "        self.cls_token = nn.Parameter(torch.zeros(1, 1, num_hiddens))\n",
        "\n",
        "        # Calculate the number of patches\n",
        "        num_patches = (img_size // patch_size) ** 2\n",
        "\n",
        "        # Learnable position embedding\n",
        "        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, num_hiddens))\n",
        "\n",
        "        # Dropout layer for the embeddings\n",
        "        self.dropout = nn.Dropout(emb_dropout)\n",
        "\n",
        "        # Stack of transformer blocks\n",
        "        self.blocks = nn.ModuleList([Block(num_hiddens, num_heads, blk_dropout, is_decoder=False) for _ in range(num_blks)])\n",
        "\n",
        "        # Layer normalization for the final representation\n",
        "        self.layer_norm = nn.LayerNorm(num_hiddens)\n",
        "\n",
        "    def forward(self, X):\n",
        "        # Convert the input image into patch embeddings\n",
        "        x = self.patch_embedding(X)\n",
        "\n",
        "        # Expand the classification token to match the batch size\n",
        "        cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)\n",
        "\n",
        "        # Concatenate the classification token with the patch embeddings\n",
        "        x = torch.cat((cls_tokens, x), dim=1)\n",
        "\n",
        "        # Add the position embedding to the patch embeddings\n",
        "        x += self.pos_embedding\n",
        "\n",
        "        # Apply dropout to the embeddings\n",
        "        x = self.dropout(x)\n",
        "\n",
        "        # Pass the embeddings through the transformer blocks\n",
        "        for block in self.blocks:\n",
        "            x = block(x)\n",
        "\n",
        "        # Apply layer normalization to the final representation\n",
        "        x = self.layer_norm(x[:, 0])\n",
        "\n",
        "        return x\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "c4c3dfa0-6e1a-42db-94b4-ea19fcdfe7a1",
          "showTitle": false,
          "title": ""
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OmTp3JN7Rw5K",
        "outputId": "4aae9670-7f98-4826-dc73-3788830b85ea"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([4, 512])"
            ]
          },
          "metadata": {},
          "execution_count": 28
        }
      ],
      "source": [
        "#For purposes of testing\n",
        "img_size, patch_size, num_hiddens, n_head, num_blks, dropout = 96, 16, 512, 8, 3, 0.1\n",
        "\n",
        "testvit = ViT(img_size, patch_size, num_hiddens, n_head, num_blks, dropout, dropout)\n",
        "vit_input = torch.zeros(batch_size, 3, img_size, img_size)\n",
        "testvit_out = testvit(vit_input)\n",
        "testvit_out.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "However, we can’t directly concatenate this to the text embeddings. We need to project this from the dimensionality of image embeddings from the vision transformer to the dimensionality of text embeddings. This is done by the vision-language projector. As mentioned before, this can be a single learnable layer followed by a non-linearity or an MLP. Here I implement an MLP for a couple of reasons.\n",
        "\n",
        "1. This is an implementation to understand how things work in a VLM. So this is more interesting than a single projection layer.\n",
        "\n",
        "2. There is an interesting current trend of keeping both the pretrained vision encoder and language decoder frozen during the VLM training phase. So giving more parameters to learn via this connection module could improve the ability of the overall VLM to generalize and help in the downstream instruction tuning process.\n",
        "\n",
        "Here’s the implementation of this projection module. It’s not too different from the MLP used in the transformer blocks.\n"
      ],
      "metadata": {
        "id": "qKq_UHuaioOC"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "c0733d1d-fb0f-4ab2-879c-9a2fd05a4dc2",
          "showTitle": false,
          "title": ""
        },
        "id": "R3z7SAf-Rw5K"
      },
      "outputs": [],
      "source": [
        "class MultiModalProjector(nn.Module):\n",
        "    def __init__(self, n_embd, image_embed_dim, dropout=0.1):\n",
        "        super().__init__()\n",
        "\n",
        "        # Define the projection network\n",
        "        self.net = nn.Sequential(\n",
        "            # Linear layer to expand the image embedding dimension\n",
        "            nn.Linear(image_embed_dim, 4 * image_embed_dim),\n",
        "\n",
        "            # GELU activation function\n",
        "            nn.GELU(),\n",
        "\n",
        "            # Linear layer to project the expanded image embeddings to the text embedding dimension\n",
        "            nn.Linear(4 * image_embed_dim, n_embd),\n",
        "\n",
        "            # Dropout layer for regularization\n",
        "            nn.Dropout(dropout)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Pass the input through the projection network\n",
        "        x = self.net(x)\n",
        "        return x\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "23de9c94-8ca3-4367-b91e-6179bd05b00d",
          "showTitle": false,
          "title": ""
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vCULLV9pRw5K",
        "outputId": "67ebf41b-fa9c-40c7-c9d6-8304bd3827c2"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([4, 128])"
            ]
          },
          "metadata": {},
          "execution_count": 30
        }
      ],
      "source": [
        "#Example values for testing\n",
        "n_embd,num_hiddens = 128, 512\n",
        "\n",
        "testmmp = MultiModalProjector(n_embd,num_hiddens)\n",
        "mmp_input = testvit_out\n",
        "testmmp_out = testmmp(mmp_input)\n",
        "testmmp_out.shape"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The final component we need to look at is the decoder language model. Here I’ve remained within the confines of the modern VLM architecture but deviated a bit in the implementation. I have integrated the projection module into the decoder model class implementation. This is because I built everything from scratch and wanted to retain the makemore causal language model architecture from Andrej Karpathy’s makemore. There’s no easy way to directly feed in reshaped embeddings in this implementation, so I’ve had to improvise a little. However in using pretrained models with the Hugging Face API or any other modern library that allows you to use pretrained large language models, you can directly feed embeddings as input to the model.\n",
        "\n",
        "That being said, what I’ve done here is an interesting exercise in that it allows you to see in pretty simple code:\n",
        "\n",
        " - How the image embeddings are reshaped using the vision language projector to match that of text embeddings.\n",
        "\n",
        " - Then concatenated with token embedding.\n",
        "\n",
        " - Subsequently combined with position embeddings and used to calculate a loss function (and eventually generate text).\n",
        "\n",
        "Essentially the text generation is conditioned on the initial image input. This can be modified in a number of ways to work with interleaved text and images, which will be useful for multi-turn conversation i.e. chat scenarios using the finetuned VLM.\n",
        "\n",
        "\n",
        "The crucial parts of this decoder implementation is given below. Note how the is_decoder flag is passed as ‘True’ to use the masked version of the self attention blocks, resulting in causal scaled dot product self attention in the language decoder. Please refer to the GitHub repo linked above for the full implementation.\n"
      ],
      "metadata": {
        "id": "vTPJqLpojC_4"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "db59fdc2-a922-4530-92f2-f51d801b3f39",
          "showTitle": false,
          "title": ""
        },
        "id": "aL9FkWjxRw5K"
      },
      "outputs": [],
      "source": [
        "class DecoderLanguageModel(nn.Module):\n",
        "    def __init__(self, n_embd, image_embed_dim, vocab_size, num_heads, n_layer, use_images=False):\n",
        "        super().__init__()\n",
        "\n",
        "        self.use_images = use_images\n",
        "\n",
        "        # Token embedding table\n",
        "        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)\n",
        "\n",
        "        # Position embedding table\n",
        "        self.position_embedding_table = nn.Embedding(1000, n_embd)\n",
        "\n",
        "        if use_images:\n",
        "            # Image projection layer to align image embeddings with text embeddings\n",
        "            self.image_projection = MultiModalProjector(n_embd, image_embed_dim)\n",
        "\n",
        "        # Stack of transformer decoder blocks\n",
        "        self.blocks = nn.Sequential(*[Block(n_embd, num_heads, is_decoder=True) for _ in range(n_layer)])\n",
        "\n",
        "        # Final layer normalization\n",
        "        self.ln_f = nn.LayerNorm(n_embd)\n",
        "\n",
        "        # Language modeling head\n",
        "        self.lm_head = nn.Linear(n_embd, vocab_size)\n",
        "\n",
        "    def forward(self, idx, image_embeds=None, targets=None):\n",
        "        # Get token embeddings from the input indices\n",
        "        tok_emb = self.token_embedding_table(idx)\n",
        "\n",
        "        if self.use_images and image_embeds is not None:\n",
        "            # Project and concatenate image embeddings with token embeddings\n",
        "            img_emb = self.image_projection(image_embeds).unsqueeze(1)\n",
        "            tok_emb = torch.cat([img_emb, tok_emb], dim=1)\n",
        "\n",
        "        # Get position embeddings\n",
        "        pos_emb = self.position_embedding_table(torch.arange(tok_emb.size(1), device=device)).unsqueeze(0)\n",
        "\n",
        "        # Add position embeddings to token embeddings\n",
        "        x = tok_emb + pos_emb\n",
        "\n",
        "        # Pass through the transformer decoder blocks\n",
        "        x = self.blocks(x)\n",
        "\n",
        "        # Apply final layer normalization\n",
        "        x = self.ln_f(x)\n",
        "\n",
        "        # Get the logits from the language modeling head\n",
        "        logits = self.lm_head(x)\n",
        "\n",
        "        if targets is not None:\n",
        "            if self.use_images and image_embeds is not None:\n",
        "                # Prepare targets by concatenating a dummy target for the image embedding\n",
        "                batch_size = idx.size(0)\n",
        "                targets = torch.cat([torch.full((batch_size, 1), -100, dtype=torch.long, device=device), targets], dim=1)\n",
        "\n",
        "            # Compute the cross-entropy loss\n",
        "            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-100)\n",
        "            return logits, loss\n",
        "\n",
        "        return logits\n",
        "\n",
        "    def generate(self, idx, image_embeds, max_new_tokens):\n",
        "        # Get the batch size and sequence length\n",
        "        B, T = idx.shape\n",
        "\n",
        "        # Initialize the generated sequence with the input indices\n",
        "        generated = idx\n",
        "\n",
        "        if self.use_images and image_embeds is not None:\n",
        "            # Project and concatenate image embeddings with token embeddings\n",
        "            img_emb = self.image_projection(image_embeds).unsqueeze(1)\n",
        "            current_output = torch.cat([img_emb, self.token_embedding_table(idx)], dim=1)\n",
        "        else:\n",
        "            current_output = self.token_embedding_table(idx)\n",
        "\n",
        "        # Generate new tokens iteratively\n",
        "        for i in range(max_new_tokens):\n",
        "            # Get the current sequence length\n",
        "            T_current = current_output.size(1)\n",
        "\n",
        "            # Get position embeddings for the current sequence length\n",
        "            current_pos_emb = self.position_embedding_table(torch.arange(T_current, device=device)).unsqueeze(0)\n",
        "\n",
        "            # Add position embeddings to the current output\n",
        "            current_output += current_pos_emb\n",
        "\n",
        "            # Pass through the transformer decoder blocks\n",
        "            for block in self.blocks:\n",
        "                current_output = block(current_output)\n",
        "\n",
        "            # Get the logits for the last token\n",
        "            logits = self.lm_head(current_output[:, -1, :])\n",
        "\n",
        "            # Apply softmax to get probabilities\n",
        "            probs = F.softmax(logits, dim=-1)\n",
        "\n",
        "            # Sample the next token based on the probabilities\n",
        "            idx_next = torch.multinomial(probs, num_samples=1)\n",
        "\n",
        "            # Concatenate the generated token to the generated sequence\n",
        "            generated = torch.cat((generated, idx_next), dim=1)\n",
        "\n",
        "            # Get the embeddings for the generated token\n",
        "            idx_next_emb = self.token_embedding_table(idx_next)\n",
        "\n",
        "            # Concatenate the generated token embeddings to the current output\n",
        "            current_output = torch.cat((current_output, idx_next_emb), dim=1)\n",
        "\n",
        "        return generated"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "b58b7a9f-34fc-45c6-8bb0-31e89a989b07",
          "showTitle": false,
          "title": ""
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dDQO0JAvRw5K",
        "outputId": "c5a0a13f-f2cd-4532-e94a-df97279566b0"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Logits shape: torch.Size([10, 51, 1000]), Loss: 7.0486555099487305\n",
            "Generated sequence shape: torch.Size([10, 70])\n"
          ]
        }
      ],
      "source": [
        "# I use n_layer to represent number of decoder transformer blocks and n_blks for the vision encoder to avoid confusion\n",
        "model = DecoderLanguageModel(n_embd=128, image_embed_dim=256, vocab_size=1000, num_heads=8, n_layer=6, use_images=True)\n",
        "model.to(device)\n",
        "# Dummy input\n",
        "B, T = 10, 50\n",
        "idx = torch.randint(0, 1000, (B, T)).to(device)\n",
        "image_embeds = torch.randn(B, 256).to(device)  # Assume image_embed_dim is 256\n",
        "\n",
        "targets = torch.randint(0, vocab_size, (B, T)).to(device)  # Only if you want to compute loss\n",
        "\n",
        "# Test forward pass\n",
        "# Check if you need to calculate loss by providing targets\n",
        "if targets is not None:\n",
        "    logits, loss = model(idx, image_embeds, targets)\n",
        "    print(f\"Logits shape: {logits.shape}, Loss: {loss}\")\n",
        "else:\n",
        "    logits = model(idx, image_embeds)  # Call without targets\n",
        "    print(f\"Logits shape: {logits.shape}\")\n",
        "\n",
        "# Test generation\n",
        "generated = model.generate(idx, image_embeds, max_new_tokens=20)\n",
        "print(f\"Generated sequence shape: {generated.shape}\")\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now that we have our three key components, we can put it all together into a Vision Language Model. The full implementation is given below. If you were to remove the assert statements for error handling, this is very simple. Coming back full circle to the outline I’ve given at the beginning of the blog, all that’s happening here is:\n",
        "\n",
        "1. Get image features from the vision encoder (Here it’s a vision transformer, but it could be any model that could generate features from an image input such as a ResNet or a traditional convolutional neural network (needless to say performance may suffer))\n",
        "\n",
        "2. A projection module for projecting image tokens to the same embedding space as text embeddings for the decoder (this projector is integrated with the decoder in this implementation)\n",
        "\n",
        "3. A decoder language model for generating text conditioned on a preceding image.\n"
      ],
      "metadata": {
        "id": "yvRS8VLOj8fo"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "6702a660-0512-4152-8d56-97f0ebd562e8",
          "showTitle": false,
          "title": ""
        },
        "id": "FjklXZrPRw5K"
      },
      "outputs": [],
      "source": [
        "class VisionLanguageModel(nn.Module):\n",
        "    def __init__(self, n_embd, image_embed_dim, vocab_size, n_layer, img_size, patch_size, num_heads, num_blks, emb_dropout, blk_dropout):\n",
        "        super().__init__()\n",
        "\n",
        "        # Set num_hiddens equal to image_embed_dim\n",
        "        num_hiddens = image_embed_dim\n",
        "\n",
        "        # Assert that num_hiddens is divisible by num_heads\n",
        "        assert num_hiddens % num_heads == 0, \"num_hiddens must be divisible by num_heads\"\n",
        "\n",
        "        # Initialize the vision encoder (ViT)\n",
        "        self.vision_encoder = ViT(img_size, patch_size, num_hiddens, num_heads, num_blks, emb_dropout, blk_dropout)\n",
        "\n",
        "        # Initialize the language model decoder (DecoderLanguageModel)\n",
        "        self.decoder = DecoderLanguageModel(n_embd, image_embed_dim, vocab_size, num_heads, n_layer, use_images=True)\n",
        "\n",
        "    def forward(self, img_array, idx, targets=None):\n",
        "        # Get the image embeddings from the vision encoder\n",
        "        image_embeds = self.vision_encoder(img_array)\n",
        "\n",
        "        # Check if the image embeddings are valid\n",
        "        if image_embeds.nelement() == 0 or image_embeds.shape[1] == 0:\n",
        "            raise ValueError(\"Something is wrong with the ViT model. It's returning an empty tensor or the embedding dimension is empty.\")\n",
        "\n",
        "        if targets is not None:\n",
        "            # If targets are provided, compute the logits and loss\n",
        "            logits, loss = self.decoder(idx, image_embeds, targets)\n",
        "            return logits, loss\n",
        "        else:\n",
        "            # If targets are not provided, compute only the logits\n",
        "            logits = self.decoder(idx, image_embeds)\n",
        "            return logits\n",
        "\n",
        "    def generate(self, img_array, idx, max_new_tokens):\n",
        "        # Get the image embeddings from the vision encoder\n",
        "        image_embeds = self.vision_encoder(img_array)\n",
        "\n",
        "        # Check if the image embeddings are valid\n",
        "        if image_embeds.nelement() == 0 or image_embeds.shape[1] == 0:\n",
        "            raise ValueError(\"Something is wrong with the ViT model. It's returning an empty tensor or the embedding dimension is empty.\")\n",
        "\n",
        "        # Generate new tokens using the language model decoder\n",
        "        generated_tokens = self.decoder.generate(idx, image_embeds, max_new_tokens)\n",
        "        return generated_tokens"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now back to where we started. The above VisualLanguageModel class neatly wraps up all the components we set out to put together.\n",
        "\n",
        "<img src=\"https://github.com/AviSoori1x/seemore/blob/main/images/vlm.png?raw=true\" width=\"500\" height=\"400\" alt=\"seemore\">"
      ],
      "metadata": {
        "id": "gniFhFaXlZmp"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "3595ca42-b3b5-4733-bd3b-0dedd59a9ee5",
          "showTitle": false,
          "title": ""
        },
        "id": "H9nFWES0Rw5L"
      },
      "outputs": [],
      "source": [
        "image_embed_dim = num_hiddens"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "2dd9be94-6ba3-4387-9a3a-139907173b4d",
          "showTitle": false,
          "title": ""
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "m-iwMzhuRw5L",
        "outputId": "ea5b603a-c759-4f69-d36f-2248e08181d0"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Output from initialization forward pass: tensor([[[-0.8307,  1.3653, -0.3604,  ...,  0.4427, -0.4578, -0.9713],\n",
            "         [-0.2743,  2.0467, -0.0566,  ..., -0.5852, -0.9578, -0.3152],\n",
            "         [-0.0215, -0.4519, -1.0820,  ...,  0.2260, -0.4066, -0.1245],\n",
            "         ...,\n",
            "         [-0.5994, -0.6862,  0.2406,  ...,  0.0324,  0.1187, -0.2532],\n",
            "         [ 0.1253,  0.2257, -0.3586,  ...,  0.1581, -0.3060, -0.0481],\n",
            "         [-0.2291, -0.4992, -0.1902,  ..., -0.1391, -1.1035, -0.1593]]],\n",
            "       device='cuda:0', grad_fn=<ViewBackward0>)\n"
          ]
        }
      ],
      "source": [
        "n_layer, block_size =  8, 32\n",
        "\n",
        "# Initialize the model\n",
        "model = VisionLanguageModel(n_embd, image_embed_dim, vocab_size,  n_layer, img_size, patch_size, n_head, num_blks, dropout, dropout)\n",
        "model.to(device)\n",
        "\n",
        "# Create dummy data with correct dimensions\n",
        "dummy_img = torch.randn(1, 3, img_size, img_size).to(device)  # Correct shape for image input\n",
        "dummy_idx = torch.randint(0, vocab_size, (1, block_size)).to(device)  # Correct shape for text input\n",
        "\n",
        "# Forward pass to initialize all parameters\n",
        "try:\n",
        "    output = model(dummy_img, dummy_idx)  # Output for debugging\n",
        "    print(\"Output from initialization forward pass:\", output)\n",
        "except RuntimeError as e:\n",
        "    print(f\"Runtime Error during forward pass: {str(e)}\")\n",
        "    print(\"Check layer configurations and input shapes.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Function to convert base64 encoded stringified images in inputs.csv in the repo to pytorch tensors so they can be used for training"
      ],
      "metadata": {
        "id": "2Z7izIhGkJDy"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "09d59b56-a8e7-4ff4-8f0f-c0edd527f47d",
          "showTitle": false,
          "title": ""
        },
        "id": "mhoM3Rt2Rw5L"
      },
      "outputs": [],
      "source": [
        "def base64_to_tensor(base64_str, img_size=96):\n",
        "    image = Image.open(io.BytesIO(base64.b64decode(base64_str)))\n",
        "    if image.mode != 'RGB':\n",
        "        image = image.convert('RGB')\n",
        "    transform = transforms.Compose([\n",
        "        transforms.Resize((img_size, img_size)),\n",
        "        transforms.ToTensor(),\n",
        "        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
        "    ])\n",
        "    return transform(image).unsqueeze(0)  # Add batch dimension"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Adjusting the dataloader in Andrej Karpathy's dataloader in makemore for the VLM training loop"
      ],
      "metadata": {
        "id": "eNVwOzS2kXwK"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "45e77024-1b6f-4b25-a946-a808305f42b8",
          "showTitle": false,
          "title": ""
        },
        "id": "qF34w9-ZRw5L"
      },
      "outputs": [],
      "source": [
        "#Adjusting the data loader from makemore for multimodal data\n",
        "def get_batch(df, batch_size, split='train', img_size=96, val_batch_size=8):\n",
        "    # Split data into training and validation sets\n",
        "    n = int(0.9 * len(df))  # first 90% will be train, rest val\n",
        "    df_train = df.iloc[:n]\n",
        "    df_val = df.iloc[n:]\n",
        "    data = df_train if split == 'train' else df_val\n",
        "    batch_size = batch_size if split == 'train' else val_batch_size\n",
        "    replace = False if split == 'train' else True\n",
        "    batch = data.sample(n=batch_size, replace=replace)\n",
        "\n",
        "    images = torch.cat([base64_to_tensor(img, img_size) for img in batch['b64string_images']], dim=0).to(device)\n",
        "    text_indices = [torch.tensor(encode(desc), dtype=torch.long) for desc in batch['caption']]\n",
        "    max_length = max(len(t) for t in text_indices)\n",
        "\n",
        "    padded_text = torch.full((batch_size, max_length), fill_value=stoi['<pad>'], dtype=torch.long).to(device)\n",
        "    for i, text in enumerate(text_indices):\n",
        "        padded_text[i, :len(text)] = text\n",
        "\n",
        "    targets = torch.cat([padded_text[:, 1:], torch.full((batch_size, 1), fill_value=stoi['<pad>'], dtype=torch.long, device=device)], dim=1)\n",
        "\n",
        "    # Truncate or pad targets to match the length of padded_text\n",
        "    if targets.size(1) > padded_text.size(1):\n",
        "        targets = targets[:, :padded_text.size(1)]\n",
        "    elif targets.size(1) < padded_text.size(1):\n",
        "        targets = torch.cat([targets, torch.full((batch_size, padded_text.size(1) - targets.size(1)), fill_value=stoi['<pad>'], dtype=torch.long, device=device)], dim=1)\n",
        "\n",
        "    return images, padded_text, targets"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Makemore training loop modified for our rudimentary Vision Language Model"
      ],
      "metadata": {
        "id": "MQKllw0pkgJF"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "605036de-0c47-4dc8-9328-1a8ad9c889f8",
          "showTitle": false,
          "title": ""
        },
        "id": "IzBrV3ZWRw5L"
      },
      "outputs": [],
      "source": [
        "#Adjusting the training loop from makemore for multimodal data\n",
        "def train_model(model, df, epochs, vocab_size, img_size=96):\n",
        "    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
        "    model.to(device)\n",
        "    for epoch in range(epochs):\n",
        "        model.train()\n",
        "        for _ in range(max_iters):\n",
        "            images, idx, targets = get_batch(df, batch_size, 'train', img_size)\n",
        "            optimizer.zero_grad()\n",
        "            logits, loss = model(images, idx, targets)\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "            if _ % eval_interval == 0:\n",
        "                print(f\"Loss at iteration {_}: {loss.item()}\")\n",
        "        val_loss = estimate_loss(model, df, 'val', img_size, val_batch_size=8)\n",
        "        print(f\"Validation Loss after epoch {epoch}: {val_loss}\")\n",
        "\n",
        "def estimate_loss(model, df, split, img_size=96, val_batch_size=8):\n",
        "    losses = []\n",
        "    model.eval()\n",
        "    for _ in range(eval_iters):\n",
        "        images, idx, targets = get_batch(df, batch_size, split, img_size, val_batch_size=val_batch_size)\n",
        "        _, loss = model(images, idx, targets)\n",
        "        losses.append(loss.item())\n",
        "    return sum(losses) / len(losses)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 39,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "1670d54d-791d-4d03-945a-12aa634af044",
          "showTitle": false,
          "title": ""
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ITWGvRmkRw5L",
        "outputId": "198364df-6a9d-4846-dc94-2ad85006e059"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(90, 2)"
            ]
          },
          "metadata": {},
          "execution_count": 39
        }
      ],
      "source": [
        "df = pd.read_csv(\"./inputs.csv\")\n",
        "#Expanding dataframe so that there's enough data to test. This is just duplicating data. A real dataset would have more rows\n",
        "df = pd.concat([df] * 30)[['b64string_images', 'caption']]\n",
        "df.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 43,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "a91f0e87-d361-40a3-a43f-51566e03810a",
          "showTitle": false,
          "title": ""
        },
        "id": "QgyBgOPJRw5L"
      },
      "outputs": [],
      "source": [
        "# display(df.head())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "3d948df0-4979-46e1-8208-629e8a95c93b",
          "showTitle": false,
          "title": ""
        },
        "id": "MFM46Cs0Rw5L"
      },
      "outputs": [],
      "source": [
        "batch_size = 16 # how many independent sequences will we process in parallel?\n",
        "block_size = 32 # what is the maximum context length for predictions?\n",
        "max_iters = 100\n",
        "eval_interval = 10\n",
        "learning_rate = 1e-3\n",
        "epochs=1\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "eval_iters = 40\n",
        "num_blks= 3\n",
        "head_size = 16\n",
        "n_embd = 128\n",
        "n_head = 8\n",
        "n_layer = 8\n",
        "dropout = 0.1\n",
        "img_size=96\n",
        "patch_size =16\n",
        "image_embed_dim = 512\n",
        "emb_dropout = blk_dropout =0.1"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's train!! Optional: Use MLFlow to track eval loss as the model is being trained"
      ],
      "metadata": {
        "id": "-Qqfkc1Tkl-f"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "metadata": {
        "application/vnd.databricks.v1+cell": {
          "cellMetadata": {
            "byteLimit": 2048000,
            "rowLimit": 10000
          },
          "inputWidgets": {},
          "nuid": "04e27881-aaee-4f40-b5cd-4e6a2de0fca7",
          "showTitle": false,
          "title": ""
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ayLj98T7Rw5L",
        "outputId": "f019ab03-6090-4cee-f2a1-1053a7e1fefd"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loss at iteration 0: 4.303300857543945\n",
            "Loss at iteration 10: 0.47614842653274536\n",
            "Loss at iteration 20: 0.09451670944690704\n",
            "Loss at iteration 30: 0.05168599635362625\n",
            "Loss at iteration 40: 0.03271578252315521\n",
            "Loss at iteration 50: 0.021071063354611397\n",
            "Loss at iteration 60: 0.012960257939994335\n",
            "Loss at iteration 70: 0.009617932140827179\n",
            "Loss at iteration 80: 0.007672513835132122\n",
            "Loss at iteration 90: 0.006618370767682791\n",
            "Validation Loss after epoch 0: 0.005623874603770673\n"
          ]
        }
      ],
      "source": [
        "# Initialize the model\n",
        "model = VisionLanguageModel(n_embd, image_embed_dim, vocab_size, n_layer, img_size, patch_size, n_head, num_blks, emb_dropout, blk_dropout)\n",
        "model.to(device)\n",
        "\n",
        "# Dummy data to initialize lazy modules\n",
        "dummy_img = torch.randn(1, 3, img_size, img_size).to(device)\n",
        "dummy_idx = torch.randint(0, vocab_size, (1, block_size)).to(device)\n",
        "model(dummy_img, dummy_idx)  # Forward pass to initialize all parameters\n",
        "\n",
        "# Train the model\n",
        "train_model(model, df, epochs, vocab_size, img_size)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Please note that in this simple example, we are training the entire system end to end, much like Kosmos-1 from Microsoft Research. I left it at this for convenience. In practice the commonly observed sequence is:\n",
        "\n",
        "1. Get pretrained vision encoder from SigLIP or CLIP (both come in difference sizes). Freeze weights (i.e. don’t update during backward pass in training)\n",
        "\n",
        "2. Get decoder only language model e.g. all the way from TinyLLaMA, Phi-2 etc. to Llama 3 (or even much bigger in the case of GPT-4 and Grok 1.5 etc.). Freeze weights.\n",
        "\n",
        "3. Implement a projection module and train a VLM module much like what we have here, but only updating the weights of this projection module. This would effectively be the pretraining phase.\n",
        "\n",
        "4. Then during the instruction finetuning keep both the projection module and the decoder language model unfrozen and update weights of both in the backward pass.\n",
        "\n",
        "I developed this on Databricks using a single T4 and MLFlow for tracking loss (during the training process). I wanted to set this up this way so that I can scale up to a GPU cluster of any size I want quite easily on Databricks, should I decide to adapt this to a more performance oriented implementation. However, you can run this anywhere, with or without a GPU. Please note that even the toy training loop with 90 samples will be painfully slow on a CPU.\n"
      ],
      "metadata": {
        "id": "k2kPV9svmUOd"
      }
    }
  ],
  "metadata": {
    "application/vnd.databricks.v1+notebook": {
      "dashboards": [],
      "language": "python",
      "notebookMetadata": {
        "pythonIndentUnit": 2
      },
      "notebookName": "seemore",
      "widgets": {}
    },
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "language_info": {
      "name": "python"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}